/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package treegenerator.arbaro;

import com.jme3.asset.AssetManager;
import com.jme3.material.Material;
import com.jme3.math.Vector2f;
import com.jme3.math.Vector3f;
import com.jme3.scene.Geometry;
import com.jme3.scene.Mesh.Mode;
import com.jme3.scene.Node;
import com.jme3.scene.VertexBuffer;
import com.jme3.scene.VertexBuffer.Format;
import com.jme3.scene.VertexBuffer.Type;
import com.jme3.scene.VertexBuffer.Usage;
import com.jme3.scene.mesh.IndexBuffer;
import com.jme3.util.BufferUtils;
import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.PrintWriter;
import java.nio.FloatBuffer;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.Random;
import net.sourceforge.arbaro.export.Progress;
import net.sourceforge.arbaro.mesh.Face;
import net.sourceforge.arbaro.mesh.LeafMesh;
import net.sourceforge.arbaro.mesh.Mesh;
import net.sourceforge.arbaro.mesh.MeshGenerator;
import net.sourceforge.arbaro.mesh.MeshGeneratorFactory;
import net.sourceforge.arbaro.mesh.MeshPart;
import net.sourceforge.arbaro.mesh.UVVector;
import net.sourceforge.arbaro.mesh.Vertex;
import net.sourceforge.arbaro.params.FloatFormat;
import net.sourceforge.arbaro.params.Params;
import net.sourceforge.arbaro.transformation.Vector;
import net.sourceforge.arbaro.tree.DefaultTreeTraversal;
import net.sourceforge.arbaro.tree.Leaf;
import net.sourceforge.arbaro.tree.Tree;
import net.sourceforge.arbaro.tree.TreeGenerator;
import net.sourceforge.arbaro.tree.TreeGeneratorFactory;

/**
 *
 * @author Jannes Plyson
 */
public class ArbaroToJMonkey{        
    MeshGenerator meshGenerator;
    AssetManager assetManager;
    ArrayList<StemTextureCreator> textureCreators;
    Params currentParams;
    
    public ArbaroToJMonkey(AssetManager assetManager) {        
        this.assetManager = assetManager;
        meshGenerator = MeshGeneratorFactory.createMeshGenerator(false);        
    }
    
    public void setTextureCreators(ArrayList<StemTextureCreator> textureCreators){
        this.textureCreators = textureCreators;
    }                
    
    public Tree createTree(InputStream is, int seed){     
        try{            
            TreeGenerator treeGenerator = TreeGeneratorFactory.createTreeGenerator();                        
            treeGenerator.readParamsFromXML(is);            
            treeGenerator.setParam("Smooth",new Double(0.5).toString());            
            treeGenerator.setSeed(seed);       
            treeGenerator.getParams();
            currentParams = treeGenerator.getParams();
            Progress progress = new Progress();
            return treeGenerator.makeTree(progress);                                    
        }catch(Exception exc){
            exc.printStackTrace();
        }
        return null;
    }
    public Tree createTree(File treeXML, int seed){        
        try{            
            InputStream is = new FileInputStream(treeXML);
            return createTree(is, seed);
        }catch(Exception exc){
            exc.printStackTrace();
        }
        return null;
    }
    
    public LeafMesh getLeafMesh(Tree tree){
        return meshGenerator.createLeafMesh(tree,false);
    }
    
    public Node createMesh(Tree tree, int treeNumber, Material trunkMaterial, Material stemMaterial, Material leafMaterial, int leafSides, long leafSeed){
        Node node = new Node("tree_" + treeNumber);
        Progress progress = new Progress();
        Mesh mesh = meshGenerator.createStemMeshByLevel(tree,progress);        
        for(int i = 0; i < tree.getLevels(); i++){
            try{                                
                ArrayList<ArrayList> lists = createLists(mesh, i);
                Geometry geo = createStem(lists, i,trunkMaterial,stemMaterial);
                node.attachChild(geo);
            }catch(Exception exc){
                exc.printStackTrace();
            }
        }        
        Geometry leafGeo = createLeaves(tree,leafMaterial,leafSides, leafSeed);
        node.attachChild(leafGeo);                
        node.updateModelBound();
        node.updateGeometricState();
        return node;
    }
    
    public Geometry createStem(ArrayList<ArrayList> lists, int level,Material trunkMaterial, Material stemMaterial){
        com.jme3.scene.Mesh mesh = new com.jme3.scene.Mesh();
        ArrayList<Vector> arbaro_vertices = (ArrayList<Vector>)lists.get(0);
        ArrayList<UVVector> arbaro_UVPoints = (ArrayList<UVVector>)lists.get(1);
        ArrayList<Face> arbaro_faces = (ArrayList<Face>)lists.get(2);
        ArrayList<Face> arbaro_uvFaces = (ArrayList<Face>)lists.get(3);                                
        ArrayList<Vector3f> vertices = new ArrayList<Vector3f>(arbaro_vertices.size()); 
        ArrayList<Vector2f> uvPoints = new ArrayList<Vector2f>(arbaro_vertices.size());
        for(int i = 0; i < arbaro_vertices.size(); i++){
            Vector v = arbaro_vertices.get(i);
            //vertices.add(new Vector3f((float)v.getX(),(float)v.getY(),(float)v.getZ()));
            vertices.add(new Vector3f((float)v.getX(),(float)v.getZ(),(float)v.getY()));
            uvPoints.add(null);
        }                        
        int[] indexes = new int[arbaro_faces.size() * 3];                
        int indexPointer = 0;        
        for(int i = 0; i < arbaro_faces.size(); i++){
            Face face = arbaro_faces.get(i);
            Face uvFace = arbaro_uvFaces.get(i);                        
            //for(int j = face.points.length-1; j >= 0; j--){  
            for(int j = 0; j < face.points.length; j++){
                //some vertices have multiple textures (first vertex in circle)                 
                int vertexIndex = (int)face.points[j] - 1;
                UVVector uv = arbaro_UVPoints.get((int)uvFace.points[j]-1);                
                Vector2f uvPoint = new Vector2f((float)uv.u, (float)uv.v);
                if(uvPoints.get(vertexIndex) == null){
                    uvPoints.set(vertexIndex, uvPoint);
                }else if(!uvPoints.get(vertexIndex).equals(uvPoint)){
                    vertices.add(new Vector3f(vertices.get(vertexIndex)));
                    uvPoints.add(new Vector2f((float)uv.u, (float)uv.v));
                    vertexIndex = vertices.size()-1;
                }
                indexes[indexPointer] = vertexIndex;                
                indexPointer++;                
            }                        
        }
        Vector3f[] positionBuffer = new Vector3f[vertices.size()];
        vertices.toArray(positionBuffer);
        Vector2f[] textureBuffer = new Vector2f[uvPoints.size()];
        uvPoints.toArray(textureBuffer);
        mesh.setBuffer(Type.Position,3, BufferUtils.createFloatBuffer(positionBuffer));                
        mesh.setBuffer(Type.TexCoord,2, BufferUtils.createFloatBuffer(textureBuffer));                
        mesh.setBuffer(Type.Index, 1, BufferUtils.createIntBuffer(indexes));
                        
        VertexBuffer[] lods = new VertexBuffer[3];
        lods[0] = new VertexBuffer(Type.Index);
        lods[0].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(indexes));        
        lods[1] = new VertexBuffer(Type.Index);
        lods[1].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(getLodIndexesStem(level, indexes, currentParams)));
        lods[2] = new VertexBuffer(Type.Index);
        lods[2].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(new int[0]));        
        mesh.setLodLevels(lods);        
        mesh.updateBound();
        Geometry geo = null;
        if(level == 0){
            geo = new Geometry("trunk",mesh);   
            geo.setMaterial(trunkMaterial);
        }else{
            geo = new Geometry("stem_" + level,mesh);        
            geo.setMaterial(stemMaterial);
        }                                
        return geo;
    }
    
    private ArrayList<ArrayList> createLists(Mesh mesh, int level){
        Enumeration parts = mesh.allParts(level);
        ArrayList<ArrayList> lists = new ArrayList<ArrayList>();
        ArrayList<Vector> vertices = new ArrayList<Vector>();
        lists.add(vertices);
        ArrayList<UVVector> uvVectors = new ArrayList<UVVector>();
        lists.add(uvVectors);
        ArrayList<Face> faces = new ArrayList<Face>();
        lists.add(faces);
        ArrayList<Face> uvFaces = new ArrayList<Face>();
        lists.add(uvFaces);
        int vertexIndex = 1;
        while(parts.hasMoreElements()){            
            MeshPart part = (MeshPart)parts.nextElement();
            Enumeration vertEnum = part.allVertices(false);
            while(vertEnum.hasMoreElements()){
                Vertex vertex = (Vertex)vertEnum.nextElement();
                vertices.add(vertex.point);            
            }
            Enumeration uvEnum = part.allVertices(true);
            while(uvEnum.hasMoreElements()){
                UVVector uvVector = (UVVector)uvEnum.nextElement();
                uvVectors.add(uvVector);          
            }
            Enumeration faceEnum = part.allFaces(mesh, vertexIndex, false);
            Enumeration uvFaceEnum = part.allFaces(mesh, 1, true);
            while(faceEnum.hasMoreElements()){
                Face face = (Face)faceEnum.nextElement();
                Face uvFace = (Face)uvFaceEnum.nextElement();
                faces.add(face);
                uvFaces.add(uvFace);
            }
            vertexIndex += part.vertexCount();
        }
        return lists;
    }                      
    
    public int[] getLodIndexesStem(int stemLevel, int indexes[], Params treeParams){
        int meshPoints = treeParams.getLevelParams(stemLevel).mesh_points;
        ArrayList<Integer> newIndexes = new ArrayList<Integer>();
        int index = 0;        
        if(meshPoints <= 3){
            return indexes;
        }else{
            int interval1 = meshPoints/3;
            int interval2 = meshPoints*2/3;
            while(index < indexes.length - meshPoints){
                if(indexes[index] == indexes[index+1] - meshPoints && indexes[index+1] == indexes[index+4] && indexes[index+2] == indexes[index+5]-meshPoints || indexes[index+2] == indexes[index+5]-2){
                    //quads 1;
                    int startIndex = indexes[index];
                    int endIndex1 = indexes[index+(meshPoints-1)*6+2];
                    int endIndex2 = indexes[index+(meshPoints-1)*6+3];
                    int endIndex3 = indexes[index+(meshPoints-1)*6+5];
                    newIndexes.add(startIndex);
                    newIndexes.add(startIndex + meshPoints);
                    newIndexes.add(startIndex + interval1);
                    
                    newIndexes.add(startIndex + interval1);
                    newIndexes.add(startIndex + meshPoints);
                    newIndexes.add(startIndex + interval1 + meshPoints);
                    
                    //quad 2
                    newIndexes.add(startIndex + interval1);
                    newIndexes.add(startIndex + interval1 + meshPoints);
                    newIndexes.add(startIndex + interval2);
                    
                    newIndexes.add(startIndex + interval2);
                    newIndexes.add(startIndex + interval1 + meshPoints);
                    newIndexes.add(startIndex + interval2 + meshPoints);
                    
                    //quad 3
                    newIndexes.add(startIndex + interval2);
                    newIndexes.add(startIndex + interval2 + meshPoints);
                    newIndexes.add(endIndex1);
                    
                    newIndexes.add(endIndex2);
                    newIndexes.add(startIndex + interval2 + meshPoints);
                    newIndexes.add(endIndex3);
                    index+= meshPoints*6;
                }else{
                    int startIndex = indexes[index];
                    int topIndex = indexes[index+1];
                    int endIndex = indexes[index+meshPoints*3-1];
                    
                    //triangle 1
                    newIndexes.add(startIndex);
                    newIndexes.add(topIndex);
                    newIndexes.add(startIndex + interval1);
                    
                    //triangle 2
                    newIndexes.add(startIndex+interval1);
                    newIndexes.add(topIndex);
                    newIndexes.add(startIndex + interval2);
                    
                    //triangle 3
                    newIndexes.add(startIndex+interval2);
                    newIndexes.add(topIndex);
                    newIndexes.add(endIndex);
                    index += meshPoints*3;                    
                }
            }
        }
        int[] returnValue = new int[newIndexes.size()];
        for(int i = 0; i < newIndexes.size(); i++){
            returnValue[i] = newIndexes.get(i);
        }        
        return returnValue;
    }
    
    
    
    public Geometry createLeaves(Tree tree, Material leafMaterial, int leafSides, long leafSeed){        
        com.jme3.scene.Mesh mesh = new com.jme3.scene.Mesh();                                       		        
        LeafMesh leafMesh = meshGenerator.createLeafMesh(tree,false);                
        int vertexOffset = 0;
        //create vertices;        
        OBJLeafVertexWriter vertexExporter = new OBJLeafVertexWriter(tree, leafMesh, vertexOffset,leafSides, leafSeed);
        tree.traverseTree(vertexExporter);
        vertexOffset = vertexExporter.leafVertexOffset;
                       
        //create faces;
        OBJLeafFaceWriter faceExporter = new OBJLeafFaceWriter(tree,leafMesh,vertexOffset);
        tree.traverseTree(faceExporter);
        vertexOffset = faceExporter.leafVertexOffset;        
        
        mesh.setBuffer(Type.Position,3, BufferUtils.createFloatBuffer(vertexExporter.getVertices()));                
        mesh.setBuffer(Type.TexCoord,2, BufferUtils.createFloatBuffer(vertexExporter.getUVPoints()));  
        int [] indexes = faceExporter.getIndexes();
        mesh.setBuffer(Type.Index, 1, BufferUtils.createIntBuffer(indexes));         
              
        VertexBuffer[] lods = new VertexBuffer[3];
        lods[0] = new VertexBuffer(Type.Index);
        lods[0].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(indexes));        
        lods[1] = new VertexBuffer(Type.Index);
        lods[1].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(getLodLeaves(indexes, leafMesh)));            
        lods[2] = new VertexBuffer(Type.Index);
        lods[2].setupData(Usage.Dynamic, 1, Format.UnsignedInt, BufferUtils.createIntBuffer(new int[0]));
        mesh.setLodLevels(lods);        
        
        mesh.updateBound();                                       
            
        Geometry geo = new Geometry("leaves",mesh);        
        geo.setMaterial(leafMaterial);                           
        return geo;
    }
    
    private com.jme3.scene.Mesh createEmitPoints(int[] indexes, OBJLeafVertexWriter vertexExporter, LeafMesh leafMesh){
        int faceCount = leafMesh.getShapeFaceCount();
        int[] result = new int[indexes.length/(faceCount*3*2)];
        Vector3f[] vertices = new Vector3f[indexes.length/(faceCount*3*2)];
        Vector2f[] uvPoints = new Vector2f[indexes.length/(faceCount*3*2)];
        int index = 0;
        for(int i = 0; i < indexes.length; i += faceCount*3*2){
            result[index] = indexes[i];  
            vertices[index] = vertexExporter.vertices.get(indexes[i]);
            uvPoints[index] = vertexExporter.uvPoints.get(indexes[i]);
            index ++;
        }        
        com.jme3.scene.Mesh emitPoints = new com.jme3.scene.Mesh();
        emitPoints.setMode(Mode.Points);
        emitPoints.setBuffer(Type.Position,3, BufferUtils.createFloatBuffer(vertices));                
        emitPoints.setBuffer(Type.TexCoord,2, BufferUtils.createFloatBuffer(uvPoints));  
        emitPoints.setBuffer(Type.Index, 1, BufferUtils.createIntBuffer(indexes));   
        return emitPoints;
    }
    
    private int[] getLodLeaves(int indexes[], LeafMesh leafMesh){
        int vertexCount = leafMesh.getShapeVertexCount();
        int faceCount = leafMesh.getShapeFaceCount();
        int[] result = new int[indexes.length/(faceCount*3)*6];
        int index = 0;
        for(int i = 0; i < indexes.length; i += faceCount*3){
            result[index] = indexes[i];
            result[index+1] = indexes[i] + vertexCount/3;
            result[index+2] = indexes[i] + vertexCount*2/3;  
            result[index+3] = indexes[i];
            result[index+4] = indexes[i] + vertexCount*2/3;
            result[index+5] = indexes[i] + vertexCount-1;            
            index += 6;
        }
        return result;
    }
    
    private void computeNormals(com.jme3.scene.Mesh mesh) {
      Vector3f vector1 = new Vector3f();
      Vector3f vector2 = new Vector3f();
      Vector3f vector3 = new Vector3f();
       
      FloatBuffer vb = mesh.getFloatBuffer(Type.Position);
      IndexBuffer ib = mesh.getIndexBuffer();
      int tCount = mesh.getTriangleCount();
      int vCount = mesh.getVertexCount();
       
      // Get the current object
      // Here we allocate all the memory we need to calculate the normals
      Vector3f[] tempNormals = new Vector3f[tCount];
      Vector3f[] normals = new Vector3f[vCount];
       
      // Go through all of the faces of this object
      for (int i = 0; i < tCount; i++) {
         BufferUtils.populateFromBuffer(vector1, vb, ib.get(i*3));
         BufferUtils.populateFromBuffer(vector2, vb, ib.get(i*3+1));
         BufferUtils.populateFromBuffer(vector3, vb, ib.get(i*3+2));
         vector1.subtractLocal(vector3);
         tempNormals[i] = vector1.cross(vector3.subtract(vector2)).normalizeLocal();
      }
       
      Vector3f sum = new Vector3f();
      int shared = 0;
       
      for (int i = 0; i < vCount; i++) {
         for (int j = 0; j < tCount; j++) {
            if (ib.get(j*3) == i
                  || ib.get(j*3+1) == i
                  || ib.get(j*3+2) == i) {
               sum.addLocal(tempNormals[j]);
                
               shared++;
            }
         }
         normals[i] = sum.divide((-shared)).normalizeLocal();
          
            sum.zero(); // Reset the sum
            shared = 0; // Reset the shared
      }
      mesh.setBuffer(Type.Normal,3,BufferUtils.createFloatBuffer(normals));
   }
}

class OBJLeafWriterBase extends DefaultTreeTraversal {
	Progress progress;
	LeafMesh leafMesh;	
	public int leafVertexOffset;
	PrintWriter w;
	long leavesProgressCount=0;
	Tree tree;

	static final NumberFormat fmt = FloatFormat.getInstance();
        
	public OBJLeafWriterBase(Tree tree,LeafMesh leafMesh,int leafVertexOffset) {
		super();				
		this.leafMesh = leafMesh;
		this.leafVertexOffset = leafVertexOffset;
	}

	public boolean enterTree(Tree tree) {
		this.tree = tree;
		return true;
	}
}
/**
 * @author wolfram
 *
 */
class OBJLeafFaceWriter extends OBJLeafWriterBase {
	ArrayList<Integer> indexes;            		
        int indexOffset;
        int offsetStep;
	/**
	 * @param pw
	 * @param leafMesh
	 * @param leafVertexOffset
	 */
	public OBJLeafFaceWriter(Tree tree,LeafMesh leafMesh,int leafVertexOffset) {
		super(tree, leafMesh, leafVertexOffset);				
                indexes = new ArrayList<Integer>();
                indexOffset = 0;
                offsetStep = leafMesh.getShapeVertexCount();
	}
	
	
	public boolean visitLeaf(Leaf l) {		                                            
		for (int i=0; i<leafMesh.getShapeFaceCount(); i++) {
                    Face face = leafMesh.shapeFaceAt(i);                                                                    
                    for(int j = 0; j < face.points.length; j++){                           
                        indexes.add(((Long)(face.points[j])).intValue()+indexOffset);                                                   
                    }                                                                                  
		}	                                
                for (int i=0; i<leafMesh.getShapeFaceCount(); i++) {
                    Face face = leafMesh.shapeFaceAt(i);                                                
                    for(int j = face.points.length-1; j >= 0; j--){
                        indexes.add(((Long)(face.points[j])).intValue()+indexOffset);
                    }                         
		}		                
                indexOffset += offsetStep;
		return true;
	}
        
        public int[] getIndexes(){            
            int[] result = new int[indexes.size()];            
            for(int i = 0; i < indexes.size(); i++){
                result[i] = indexes.get(i);
            }
            return result;
        }

}

/**
 * @author wolfram
 *
 */
class OBJLeafVertexWriter extends OBJLeafWriterBase {
	String type;
	long vertexProgressCount=0;
        ArrayList<Vector3f> vertices = new ArrayList<Vector3f>(); 	
        ArrayList<Vector2f> uvPoints = new ArrayList<Vector2f>();
        int leafSides;
        Random random;
	/**
	 * @param pw
	 * @param leafMesh
	 * @param leafVertexOffset
	 */
	public OBJLeafVertexWriter(Tree tree, LeafMesh leafMesh,int leafVertexOffset, int leafSides, long leafSeed) {
		super(tree, leafMesh, leafVertexOffset);
		this.type=type;
                this.leafSides = leafSides;
                random = new Random(leafSeed);
	}

	public boolean visitLeaf(Leaf l) {
                double uOffset = ((double)random.nextInt(leafSides))/leafSides;
                double vOffset = ((double)random.nextInt(leafSides))/leafSides;
                
		for (int i=0; i<leafMesh.getShapeVertexCount(); i++) {			                    
                    Vector v = l.getTransformation().apply(leafMesh.shapeVertexAt(i).point);                                        
                    //vertices.add(new Vector3f((float)v.getX(),(float)v.getY(),(float)v.getZ()));
                    vertices.add(new Vector3f((float)v.getX(),(float)v.getZ(),(float)v.getY()));
                    UVVector uv = leafMesh.shapeUVAt(i);
                    uvPoints.add(new Vector2f((float)(uv.u/leafSides+uOffset),(float)(uv.v/leafSides+vOffset)));
		}				
		return true;
	}	
        
        public Vector3f[] getVertices(){
            Vector3f[] result = new Vector3f[vertices.size()];
            return vertices.toArray(result);
        }
        
        public Vector2f[] getUVPoints(){
            Vector2f[] result = new Vector2f[uvPoints.size()];
            return uvPoints.toArray(result);
        }
}